package com.nulldev.util.networking.misc;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.math.BigInteger;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;

import com.nulldev.util.JVM.Arguments;
import com.nulldev.util.JVM.JVM;
import com.nulldev.util.JVM.reflections.Reflections;
import com.nulldev.util.VariableAPI.ArrayUtils;
import com.nulldev.util.VariableAPI.StringsUtil;
import com.nulldev.util.VariableAPI.util.random.GlobalRandom;
import com.nulldev.util.data.Variables;
import com.nulldev.util.data.Base64.Base64;
import com.nulldev.util.io.file.FileUtil;
import com.nulldev.util.logging.Logger;
import com.nulldev.util.logging.LoggerFactory;
import com.nulldev.util.manage.MarkedForRemoval;
import com.nulldev.util.manage.nullUtil;

public class TLSUtil {

	private static boolean __jvm_init = false;
	static {
		if (!__jvm_init) {
			try {
				/* Enable all crypto, ignore all export restrictions. */
				Security.setProperty("crypto.policy", "unlimited");
				/* Use legacy compatible close_notify logic. */
				System.setProperty("jdk.tls.acknowledgeCloseNotify", "true");
				/* Enable TLS session reuse */
				System.setProperty("jdk.tls.client.enableSessionTicketExtension", "true");
				System.setProperty("jdk.tls.server.enableSessionTicketExtension", "true");
				/* Lower the EDHE key size. */
				System.setProperty("jdk.tls.ephemeralDHKeySize", "matched");
				/* Enable the MFLN extension */
				System.setProperty("jsse.enableMFLNExtension", "true");
				/* Disable TLS client re-negotiation. */
				System.setProperty("jdk.tls.rejectClientInitiatedRenegotiation", "true");
			} catch (Exception ex) {
				ex.printStackTrace();
			}
			try {
				if (Arguments.hasArgument("--TLSUtil.enableBouncyCastle") && !nullUtil.IS_ANDROID_BUILD) {
					if (Reflections.is("com.nulldev.util.internal.externals.bouncycastle.jsse.provider.BouncyCastleJsseProvider")
							&& Reflections.is("com.nulldev.util.internal.externals.bouncycastle.jce.provider.BouncyCastleProvider")
							&& Reflections.is("com.nulldev.util.internal.externals.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider")) {
						Security.insertProviderAt(
								(Provider) Reflections.make("com.nulldev.util.internal.externals.bouncycastle.jce.provider.BouncyCastleProvider"), 1);
						Security.insertProviderAt(
								(Provider) Reflections.make("com.nulldev.util.internal.externals.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider"), 2);
						Security.insertProviderAt(
								(Provider) Reflections.make("com.nulldev.util.internal.externals.bouncycastle.jsse.provider.BouncyCastleJsseProvider"), 3);
					} else if (Reflections.is("org.bouncycastle.jsse.provider.BouncyCastleJsseProvider")
							&& Reflections.is("org.bouncycastle.jce.provider.BouncyCastleProvider")
							&& Reflections.is("org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider")) {
						Security.insertProviderAt((Provider) Reflections.make("org.bouncycastle.jce.provider.BouncyCastleProvider"), 1);
						Security.insertProviderAt((Provider) Reflections.make("org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider"), 2);
						Security.insertProviderAt((Provider) Reflections.make("org.bouncycastle.jsse.provider.BouncyCastleJsseProvider"), 3);
					}
				}
			} catch (Exception ex) {
				ex.printStackTrace();
			}
			__jvm_init = true;
		}
	}

	public static void __static_init() {
		/* NO-OP */
	}

	public static enum TLSVersion {
		/* SSL (deprecated in nullUtil 3.x) */
		@Deprecated
		SSL2(0.0),
		@Deprecated
		SSL3(0.1),
		/* TLS */
		TLS_1_0(1.0),
		TLS_1_1(1.1),
		TLS_1_2(1.2),
		TLS_1_3(1.3),

		UNKNOWN(-1);

		private final double versionID;

		private TLSVersion(final double ver) {
			this.versionID = ver;
		}

		public double getVersion() {
			return this.versionID;
		}
	}

	public static boolean hasInstalledProviders() {
		if (nullUtil.IS_ANDROID_BUILD) {
			return true;
		} else {
			return false;
		}
	}

	public static boolean supportsTLS(final TLSVersion tls) {
		if (tls == null || tls == TLSVersion.UNKNOWN)
			return false;
		return getMaxSupportedVersion().versionID >= tls.versionID;
	}

	public static TLSVersion getMaxSupportedVersion() {
		try {
			final String[] protocols = SSLContext.getDefault().getSupportedSSLParameters().getProtocols();
			final String e = ArrayUtils.getOrFinal(protocols, "TLSv1.3", "TLSv1.2", "TLSv1.1", "TLSv1");
			final String v = e.substring(e.indexOf("v") + 1);
			final double f = Double.parseDouble(v);
			for (final TLSVersion tv : TLSVersion.values()) {
				if (tv.versionID == f)
					return tv;
			}
		} catch (Exception e) {
			final float version = JVM.version();
			if (version >= 11f) {
				return TLSVersion.TLS_1_3;
			} else if (version >= 1.7f) {
				return TLSVersion.TLS_1_2;
			} else if (version >= 1.6f) {
				return TLSVersion.TLS_1_1;
			} else {
				return TLSVersion.TLS_1_0;
			}
		}
		return TLSVersion.UNKNOWN;
	}

	public static void main(String[] args) {
		JVM.println("Running: " + JVM.getFullVersion() + " (" + JVM.uncachedVersionString() + ") ");
		JVM.println("This environment supports: " + getMaxSupportedVersion());
		JVM.println("This environment's TLS implementation: " + getTLSImplementation());
		JVM.println("This environment supports TLS 1.3: " + allowTLS13());
	}

	public static String getTLSForSSL() {
		final TLSVersion v = getMaxSupportedVersion();
		if (v == TLSVersion.TLS_1_3 && !allowTLS13())
			return "TLSv1.2";
		else
			return "TLSv" + v.versionID;
	}

	public static TLSVersion getSSLTLS() {
		final TLSVersion v = getMaxSupportedVersion();
		/* The more permanent broken TLSv1.3 workaround */
		if (v == TLSVersion.TLS_1_3 && !allowTLS13())
			return TLSVersion.TLS_1_2;
		else
			return v;
	}

	public static String getTLSVersion() {
		try {
			final String[] protocols = SSLContext.getDefault().getSupportedSSLParameters().getProtocols();
			return ArrayUtils.getOrFinal(protocols, "TLSv1.3", "TLSv1.2", "TLSv1.1", "TLSv1");
		} catch (NoSuchAlgorithmException e) {
			final float version = JVM.version();
			if (version >= 11f) {
				return "TLSv1.3";
			} else if (version >= 1.7f) {
				return "TLSv1.2";
			} else if (version >= 1.6f) {
				return "TLSv1.1";
			} else {
				return "TLSv1";
			}
		}
	}

	public static boolean allowTLS13() {
		/* Android has it's own TLS 1.3 third-party implementations. */
		if (nullUtil.IS_ANDROID_BUILD)
			return true;
		/* Java 7 and older don't support TLS 1.3 */
		/* Technically we don't support older versions of Java... */
		if (JVM.version() < 1.8f)
			return false;
		/* All TLS 1.3 bugs were ironed out from 11.0.8 and onwards. */
		if (JVM.getFullVersion().newerThan(11, 0, 8))
			return true;
		/* 8u352 enables TLS 1.3 for Client (was enabled for Server earlier) */
		if (JVM.version() == 1.8f && JVM.getFullVersion().newerThan(1, 8, 352))
			return true;
		return getTLSImplementation() == TLSImpl.OpenJSSE || getTLSImplementation() == TLSImpl.BouncyCastle;
	}

	public static enum TLSImpl {
		/* Core implementations */

		UNKNOWN,
		JRE_NATIVE,

		/* Third party implementations */

		OpenJSSE,
		BouncyCastle,

		/* nullUtil Implementations */

		@Deprecated
		@MarkedForRemoval(reason = "Removed in nullUtil 3.0.2, build 342.")
		TLS12_BACKPORT,
		MICROTLS,
	}

	private static final Logger log = LoggerFactory.getLogger();

	public static TLSImpl getTLSImplementation() {
		return getTLSImplementation(false);
	}

	public static TLSImpl getTLSImplementation(final boolean useDefault) {
		try {
			final SSLContext ssl = (useDefault ? SSLContext.getDefault() : SSLContext.getInstance(getTLSVersion()));
			final Class<?> sslc = ssl.getProvider().getClass();
			switch (sslc.getName()) {
				case "com.sun.net.ssl.internal.ssl.Provider":
					return TLSImpl.JRE_NATIVE;
				case "com.nulldev.util.internal.backport.tls12.Tls12Provider":
					return TLSImpl.TLS12_BACKPORT;
				case "com.nulldev.util.internal_projects.microTLS.provider.MicroTLS":
					return TLSImpl.MICROTLS;
				case "org.bouncycastle.jsse.provider.BouncyCastleJsseProvider":
				case "com.nulldev.util.internal.externals.bouncycastle.jsse.provider.BouncyCastleJsseProvider":
					return TLSImpl.BouncyCastle;
				case "org.openjsse.net.ssl.OpenJSSE":
					return TLSImpl.OpenJSSE;
				default: {
					log.warn("Unknown implementation class: " + sslc.getName());
					return TLSImpl.UNKNOWN;
				}
			}
		} catch (Throwable r) {
			JVM.errorln(r);
			return TLSImpl.UNKNOWN;
		}
	}

	public static void tryTLSInit() {
	}

	public static SSLSocketFactory getSocketFactoryPEM(final String pemPath, final boolean privateIsEC) throws Exception {
		final SSLContext context = SSLContext.getInstance("TLS");

		final byte[] certAndKey = FileUtil.readFileToBytes(new File(pemPath));

		final byte[] certBytes = parseDERFromPEM(certAndKey, "-----BEGIN CERTIFICATE-----", "-----END CERTIFICATE-----");
		final byte[] keyBytes = parseDERFromPEM(certAndKey, "-----BEGIN PRIVATE KEY-----", "-----END PRIVATE KEY-----");

		final X509Certificate cert = generateCertificateFromDER(certBytes);
		final Key key = privateIsEC ? generateECPrivateKeyFromDER(keyBytes) : generateRSAPrivateKeyFromDER(keyBytes);

		final KeyStore keystore = KeyStore.getInstance("JKS");
		keystore.load(null);
		keystore.setCertificateEntry("cert-alias", cert);
		keystore.setKeyEntry("key-alias", key, "changeit".toCharArray(), new Certificate[]
			{ cert });

		final KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
		kmf.init(keystore, "changeit".toCharArray());

		final KeyManager[] km = kmf.getKeyManagers();

		context.init(km, null, GlobalRandom.getGlobalInsecureRandom());

		return context.getSocketFactory();
	}

	public static KeyManager[] getKeyManagers(final String certPemPath, final String keyPemPath, final boolean privateIsEC) throws Exception {
		final byte[] certData = FileUtil.readFileToBytes(new File(certPemPath));
		final byte[] keyData = FileUtil.readFileToBytes(new File(keyPemPath));

		final int certCount = getAmountOfCerts(certData);
		Variables.checkArgument(certCount != 0);

		final byte[] keyBytes = parseDERFromPEM(keyData, "-----BEGIN PRIVATE KEY-----", "-----END PRIVATE KEY-----");
		final Key key = privateIsEC ? generateECPrivateKeyFromDER(keyBytes) : generateRSAPrivateKeyFromDER(keyBytes);

		if (certCount == 1) {
			final byte[] certBytes = parseDERFromPEM(certData, "-----BEGIN CERTIFICATE-----", "-----END CERTIFICATE-----");

			final X509Certificate cert = generateCertificateFromDER(certBytes);

			final KeyStore keystore = KeyStore.getInstance("JKS");
			keystore.load(null);
			keystore.setCertificateEntry("cert-alias", cert);
			keystore.setKeyEntry("key-alias", key, "changeit".toCharArray(), new Certificate[]
				{ cert });

			final KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
			kmf.init(keystore, "changeit".toCharArray());

			final KeyManager[] km = kmf.getKeyManagers();
			return km;
		} else {
			final Certificate[] certArray = new Certificate[certCount];
			final byte[][] certBytes = parseAllDERFromPEM(certData, "-----BEGIN CERTIFICATE-----", "-----END CERTIFICATE-----");
			for (int i = 0; i < certCount; i++) {
				final X509Certificate cert = generateCertificateFromDER(certBytes[i]);
				certArray[i] = cert;
			}

			final KeyStore keystore = KeyStore.getInstance("JKS");
			keystore.load(null);
			for (int i = 0; i < certCount; i++) {
				keystore.setCertificateEntry("cert-alias-" + i, certArray[i]);
			}
			keystore.setKeyEntry("key-alias", key, "changeit".toCharArray(), certArray);

			final KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
			kmf.init(keystore, "changeit".toCharArray());

			final KeyManager[] km = kmf.getKeyManagers();
			return km;
		}
	}

	private static int getAmountOfCerts(final byte[] certData) {
		final String n = new String(certData);
		final int certStart = StringsUtil.countMatches(n, "-----BEGIN CERTIFICATE-----");
		final int certEnd = StringsUtil.countMatches(n, "-----END CERTIFICATE-----");
		if (certStart == certEnd) {
			return certStart;
		} else {
			log.warn("[getAmountOfCerts] Certificate count mismatch! [s=" + certStart + ",e=" + certEnd + "]");
			return (certStart + certEnd) / 2;
		}
	}

	public static KeyManager[] getKeyManagers(final String pemPath, final boolean privateIsEC) throws Exception {
		final byte[] certAndKey = FileUtil.readFileToBytes(new File(pemPath));

		final int certCount = getAmountOfCerts(certAndKey);
		Variables.checkArgument(certCount != 0);

		final byte[] keyBytes;
		final boolean usePKCS1;

		if (new String(certAndKey).contains("-----BEGIN RSA PRIVATE KEY-----")) {
			keyBytes = parseDERFromPEM(certAndKey, "-----BEGIN RSA PRIVATE KEY-----", "-----END RSA PRIVATE KEY-----");
			usePKCS1 = true;
		} else {
			keyBytes = parseDERFromPEM(certAndKey, "-----BEGIN PRIVATE KEY-----", "-----END PRIVATE KEY-----");
			usePKCS1 = false;
		}

		final Key key;

		if (privateIsEC) {
			key = generateECPrivateKeyFromDER(keyBytes);
		} else if (usePKCS1) {
			key = generateRSAPrivateKeyFromDER1(keyBytes);
		} else {
			key = generateRSAPrivateKeyFromDER(keyBytes);
		}

		if (certCount == 1) {
			final byte[] certBytes = parseDERFromPEM(certAndKey, "-----BEGIN CERTIFICATE-----", "-----END CERTIFICATE-----");

			final X509Certificate cert = generateCertificateFromDER(certBytes);

			final KeyStore keystore = KeyStore.getInstance("JKS");
			keystore.load(null);
			keystore.setCertificateEntry("cert-alias", cert);
			keystore.setKeyEntry("key-alias", key, "changeit".toCharArray(), new Certificate[]
				{ cert });

			final KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
			kmf.init(keystore, "changeit".toCharArray());

			final KeyManager[] km = kmf.getKeyManagers();
			return km;
		} else {
			final Certificate[] certArray = new Certificate[certCount];
			final byte[][] certBytes = parseAllDERFromPEM(certAndKey, "-----BEGIN CERTIFICATE-----", "-----END CERTIFICATE-----");
			for (int i = 0; i < certCount; i++) {
				final X509Certificate cert = generateCertificateFromDER(certBytes[i]);
				certArray[i] = cert;
			}

			final KeyStore keystore = KeyStore.getInstance("JKS");
			keystore.load(null);
			for (int i = 0; i < certCount; i++) {
				keystore.setCertificateEntry("cert-alias-" + i, certArray[i]);
			}
			keystore.setKeyEntry("key-alias", key, "changeit".toCharArray(), certArray);

			final KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
			kmf.init(keystore, "changeit".toCharArray());

			final KeyManager[] km = kmf.getKeyManagers();
			return km;
		}
	}

	protected static byte[] parseDERFromPEM(final byte[] pem, final String beginDelimiter, final String endDelimiter) {
		final String data = new String(pem);
		String[] tokens = data.split(beginDelimiter);
		tokens = tokens[1].split(endDelimiter);
		return Base64.decodeBase64(tokens[0]);
	}

	protected static byte[][] parseAllDERFromPEM(final byte[] pem, final String beginDelimiter, final String endDelimiter) {
		final String data = new String(pem);
		final String[] _tokens = data.replaceAll("\\n", "").split(beginDelimiter);
		final String[] tokens = new String[_tokens.length - 1];
		System.arraycopy(_tokens, 1, tokens, 0, tokens.length);
		final byte[][] output = new byte[tokens.length][];
		for (int i = 0; i < tokens.length; i++) {
			final String token = tokens[i];
			final String[] split = token.split(endDelimiter);
			output[i] = Base64.decodeBase64(split[0]);
		}
		return output;
	}

	protected static ECPrivateKey generateECPrivateKeyFromDER(final byte[] keyBytes) throws InvalidKeySpecException, NoSuchAlgorithmException {
		final PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);

		final KeyFactory factory = KeyFactory.getInstance("EC");
		return (ECPrivateKey) factory.generatePrivate(spec);
	}

	private static byte[] buildPkcs8KeyFromPkcs1Key(final byte[] innerKey) {
		final byte[] result = new byte[innerKey.length + 26];
		System.arraycopy(Base64.instance.decode("MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKY="), 0, result, 0, 26);
		System.arraycopy(BigInteger.valueOf(result.length - 4).toByteArray(), 0, result, 2, 2);
		System.arraycopy(BigInteger.valueOf(innerKey.length).toByteArray(), 0, result, 24, 2);
		System.arraycopy(innerKey, 0, result, 26, innerKey.length);
		return result;
	}

	protected static RSAPrivateKey generateRSAPrivateKeyFromDER1(final byte[] _keyBytes) throws InvalidKeySpecException, NoSuchAlgorithmException {
		final byte[] keyBytes = buildPkcs8KeyFromPkcs1Key(_keyBytes);
		return generateRSAPrivateKeyFromDER(keyBytes);
	}

	protected static RSAPrivateKey generateRSAPrivateKeyFromDER(final byte[] keyBytes) throws InvalidKeySpecException, NoSuchAlgorithmException {
		final PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);

		final KeyFactory factory = KeyFactory.getInstance("RSA");
		return (RSAPrivateKey) factory.generatePrivate(spec);
	}

	protected static X509Certificate generateCertificateFromDER(final byte[] certBytes) throws CertificateException {
		final CertificateFactory factory = CertificateFactory.getInstance("X.509");
		return (X509Certificate) factory.generateCertificate(new ByteArrayInputStream(certBytes));
	}
}
